import torch
import torch.nn as nn
from utils import init
from torch.distributions.distribution import Distribution
from torch.distributions import constraints
from torch.distributions.utils import probs_to_logits, logits_to_probs, lazy_property

"""
 PyTorch distributions have bugs, modify it !!!
"""
class myCategorical(Distribution):
    arg_constraints = {'probs': constraints.simplex,
                       'logits': constraints.real}
    has_enumerate_support = True

    def __init__(self, probs=None, logits=None, validate_args=None):
        if (probs is None) == (logits is None):
            raise ValueError("Either `probs` or `logits` must be specified, but not both.")
        if probs is not None:
            if probs.dim() < 1:
                raise ValueError("`probs` parameter must be at least one-dimensional.")
            self.probs = probs / probs.sum(-1, keepdim=True)
        else:
            if logits.dim() < 1:
                raise ValueError("`logits` parameter must be at least one-dimensional.")

            # Normalize, 
            # self.logits are real logits
            # which is self.logits = [ log(p1), log(p2), log(p3), ... ], 
            # while input logits are from (-inf, inf), input_logits = [ log(q1), log(q2), log(q3), ... ]
            # the logsumexp make a convertion elegantly: p1 = q1/(q1+q2+...), p2 = q2/(q1+q2+...)

            self.logits = logits - logits.logsumexp(dim=-1, keepdim=True)   # exp -> sum exp -> log(sum(exp()))
        self._param = self.probs if probs is not None else self.logits
        self._num_events = self._param.size()[-1]
        batch_shape = self._param.size()[:-1] if self._param.ndimension() > 1 else torch.Size()
        super(myCategorical, self).__init__(batch_shape, validate_args=validate_args)

    def expand(self, batch_shape, _instance=None):
        new = self._get_checked_instance(Categorical, _instance)
        batch_shape = torch.Size(batch_shape)
        param_shape = batch_shape + torch.Size((self._num_events,))
        if 'probs' in self.__dict__:
            new.probs = self.probs.expand(param_shape)
            new._param = new.probs
        if 'logits' in self.__dict__:
            new.logits = self.logits.expand(param_shape)
            new._param = new.logits
        new._num_events = self._num_events
        super(Categorical, new).__init__(batch_shape, validate_args=False)
        new._validate_args = self._validate_args
        return new

    def _new(self, *args, **kwargs):
        return self._param.new(*args, **kwargs)

    @constraints.dependent_property
    def support(self):
        return constraints.integer_interval(0, self._num_events - 1)

    @lazy_property
    def logits(self):   
        return probs_to_logits(self.probs)

    @lazy_property
    def probs(self):    # whoooa 可以看到，logits虽然是定义成一个方法的形式，但是加上@property后，可以直接执行，当成属性访问， 这里迷惑了vs的调试器
        return logits_to_probs(self.logits)

    @property
    def param_shape(self):
        return self._param.size()

    @property
    def mean(self):
        return torch.full(self._extended_shape(), nan, dtype=self.probs.dtype, device=self.probs.device)

    @property
    def variance(self):
        return torch.full(self._extended_shape(), nan, dtype=self.probs.dtype, device=self.probs.device)
        
    def convert_shit_fashion(self,x):
        thread_num = x.shape[0]
        agent_num = x.shape[1]
        last_dim = x.shape[2]

        return x.transpose(0,1).contiguous().view(-1,last_dim), thread_num, agent_num, last_dim
        
    def convert_normal_fashion(self,x, thread_num, agent_num, last_dim):
        return x.view(agent_num, thread_num, last_dim).transpose(0,1)

    def sample_raw(self, sample_shape=torch.Size()):
        if not isinstance(sample_shape, torch.Size):
            sample_shape = torch.Size(sample_shape)
        #probs_2d = self.probs.reshape(-1, self._num_events)
        tmp_probs_2d, thread_num, agent_num, last_dim = self.convert_shit_fashion(self.probs)
        samples_2d = torch.multinomial(tmp_probs_2d, sample_shape.numel(), True).T
        out = self.convert_normal_fashion(samples_2d, thread_num, agent_num, 1)
        return out#samples_2d.reshape(self._extended_shape(sample_shape))

    def mode(self):
        return self.probs.argmax(dim=-1, keepdim=True)


    def sample(self):
        res = self.sample_raw()
        return res

    def log_prob(self, value):
        if self._validate_args:
            self._validate_sample(value)
        value = value.long().unsqueeze(-1)
        value, log_pmf = torch.broadcast_tensors(value, self.logits)
        value = value[..., :1]
        return log_pmf.gather(-1, value).squeeze(-1)

    def log_probs(self, actions):
        s_actions = actions.squeeze(-1)
        res = self.log_prob(s_actions)
        return res.unsqueeze(-1)

    def entropy(self):
        min_real = torch.finfo(self.logits.dtype).min
        logits = torch.clamp(self.logits, min=min_real)
        p_log_p = logits * self.probs
        return -p_log_p.sum(-1)

    def enumerate_support(self, expand=True):
        num_events = self._num_events
        values = torch.arange(num_events, dtype=torch.long, device=self._param.device)
        values = values.view((-1,) + (1,) * len(self._batch_shape))
        if expand:
            values = values.expand((-1,) + self._batch_shape)
        return values

class MyCategorical(nn.Module):
    def __init__(self, num_inputs, num_outputs):
        super(MyCategorical, self).__init__()

        init_ = lambda m: init(m,
              nn.init.orthogonal_,
              lambda x: nn.init.constant_(x, 0),
              gain=0.01)

        self.linear = init_(nn.Linear(num_inputs, num_outputs))

    def forward(self, x):
        x = self.linear(x)
        return myCategorical(logits=x)




